//
// Created by wanghaiwei on 2019/12/24.
//
#pragma once

#include <string>
#include <vector>
#include <boost/serialization/string.hpp>
#include <boost/serialization/vector.hpp>

#ifndef LIBDL_CORE_MODULE_STRUCTS_H
#define LIBDL_CORE_MODULE_STRUCTS_H

class module_structs {
    friend class boost::serialization::access;

    template<class Archive>
    void serialize(Archive &archive, const unsigned int version) {
        archive & submodules;
        archive & buffers;
        archive & parameters;
        archive & name;
    }

public:
    class tensor_struct {
        friend class boost::serialization::access;

        template<class Archive>
        void serialize(Archive &archive, const unsigned int version) {
            archive & name;
            archive & size;
        }

    public:
        std::string name;
        std::vector<int64_t> size;

        tensor_struct(const std::string &name, c10::IntArrayRef size) {
            this->name = name;
            this->size = size.vec();
        }

        tensor_struct() = default;

        c10::IntArrayRef get_sizes() {
            c10::IntArrayRef array(size);
            return array;
        }
    };

public:
    std::vector<module_structs> submodules;
    std::vector<tensor_struct> parameters;
    std::vector<tensor_struct> buffers;
    std::string name;

    bool contains(const std::string &module_name) {
        for (const auto& item : this->submodules) {
            if (item.name == module_name)
                return true;
        }
        return false;
    }
};

#endif //LIBDL_CORE_MODULE_STRUCTS_H
